import os
import re
import json

from typing import Any, Dict
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser


class Promt_Manager:
    def __init__(self, cfg: Dict[str, Any]):
        self.file_path = cfg.get('promt_path')
        self.retry = cfg.get('retry', 3)

        model = ChatOpenAI(
            model=cfg.get('promt_model'),
            base_url=cfg.get('base_url'),
            api_key=cfg.get('api_key'),
            temperature=0.0
        )
        
        prompt = ChatPromptTemplate.from_messages([
            (
                "system",
                "当前 promt.txt 文件内容如下：\n"
                "{promt_lines}\n"
                "用户最新输入：{user_input}\n\n"
                "请判断用户是否明确要求需要改变说话方式"
                "如果有则在**添加、修改或删除**中选一条最合适的 prompt。\n"
                "如需添加，输出 JSON：\n"
                "{{\"action\": \"add\", \"new_line\": \"要添加的新内容\"}}\n"
                "如需修改，输出：\n"
                "{{\"action\": \"update\", \"line_number\": 行号, \"new_line\": \"新的内容\"}}\n"
                "如需删除，输出：\n"
                "{{\"action\": \"delete\", \"line_number\": 行号}}\n"
                "如无需修改，只返回：\n"
                "{{\"action\": \"none\"}}\n"
                "使用你认为好的提示词，不必完全照搬用户原话"
                "注意：**只有用户说得很明确时才操作，不要猜测。**\n"
                "不要输出多余的内容，只返回标准 JSON！"
            ),
            ("user", "{user_input}")
        ])
        parser = StrOutputParser()
        self.chain = prompt | model | parser

    def get_all(self) -> Dict[int, str]:
        lines = {}
        if not os.path.exists(self.file_path):
            return lines
        with open(self.file_path, "r", encoding="utf-8") as f:
            for i, line in enumerate(f, 1):
                lines[i] = line.rstrip('\n')
        return lines
    
    def read_promt_lines(self):
        lines = []
        if os.path.exists(self.file_path):
            with open(self.file_path, "r", encoding="utf-8") as f:
                for idx, line in enumerate(f, 1):
                    lines.append(f"{idx}. {line.rstrip()}")
        return "\n".join(lines) if lines else "(当前为空)"

    def update_line(self, line_number: int, new_line: str) -> bool:
        if not os.path.exists(self.file_path):
            return False
        with open(self.file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()
        if 1 <= line_number <= len(lines):
            lines[line_number - 1] = new_line.rstrip('\n') + '\n'
            with open(self.file_path, "w", encoding="utf-8") as f:
                f.writelines(lines)
            return True
        return False

    def delete_line(self, line_number: int) -> bool:
        if not os.path.exists(self.file_path):
            return False
        with open(self.file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()
        if 1 <= line_number <= len(lines):
            del lines[line_number - 1]
            with open(self.file_path, "w", encoding="utf-8") as f:
                f.writelines(lines)
            return True
        return False

    def add_line(self, new_line: str) -> bool:
        """追加一行"""
        try:
            with open(self.file_path, "a", encoding="utf-8") as f:
                f.write(new_line.rstrip('\n') + '\n')
            return True
        except Exception as e:
            print(f"添加失败: {e}")
            return False

    def route(self, user_input: str) -> None:
        promt_lines = self.read_promt_lines()
        for i in range(self.retry):
            raw = self.chain.invoke({"promt_lines": promt_lines, "user_input": user_input})
            print(raw)
            if raw.strip().startswith("```"):
                json_str = re.sub(r"^```(?:json)?", "", raw.strip())
                json_str = json_str.rstrip("`").strip()
            else:
                json_str = raw.strip()
            json_str = re.sub(r",\s*([\]}])", r"\1", json_str)
            try:
                resp = json.loads(json_str)
                action = resp.get("action")
                if action == "add":
                    self.add_line(resp["new_line"])
                elif action == "update":
                    self.update_line(resp["line_number"], resp["new_line"])
                elif action == "delete":
                    self.delete_line(resp["line_number"])
                elif action == "none":
                    pass
                else:
                    print("未知action")
                break
            except Exception as e:
                print(f"第 {i} 次调用解析失败，返回内容：{raw}")
        else:
            print("LLM 判断未成功。")
